#include <opencv2/opencv.hpp> // 包含OpenCV核心头文件
#include <vector>
#include <iostream>
using namespace cv;
using namespace std;
const Scalar LOWER_BLUE = Scalar(80, 20, 40);    // 蓝色HSV下限
const Scalar UPPER_BLUE = Scalar(140, 255, 255); // 蓝色HSV上限
const double AREA_THRESHOLD = 50;                // 灯条最小面积
const double ASPECT_RATIO = 3.0;                 // 灯条长宽比

Mat detectLightBars(Mat frame)
{
    Mat hsv, mask, result;
    frame.copyTo(result);
    cvtColor(frame, hsv, COLOR_BGR2HSV);
    // 颜色提取
    inRange(hsv, LOWER_BLUE, UPPER_BLUE, mask);
    // 去噪填充
    Mat kernel = getStructuringElement(MORPH_RECT, Size(3, 3));
    morphologyEx(mask, mask, MORPH_OPEN, kernel);
    morphologyEx(mask, mask, MORPH_CLOSE, kernel);
    // 找轮廓
    vector<vector<Point>> contours;
    vector<Vec4i> hierarchy;
    findContours(mask, contours, hierarchy, RETR_EXTERNAL, CHAIN_APPROX_SIMPLE);
    for (size_t i = 0; i < contours.size(); i++)
    {
        // 筛选
        double contourArea = cv::contourArea(contours[i]);
        if (contourArea < AREA_THRESHOLD)
        {
            continue;
        }
        RotatedRect rotatedRect = minAreaRect(contours[i]);
        Size2f rectSize = rotatedRect.size;
        double ratio = max(rectSize.width, rectSize.height) / min(rectSize.width, rectSize.height);
        if (ratio < ASPECT_RATIO)
        {
            continue;
        }
        // 标记
        Point2f rectPoints[4];
        rotatedRect.points(rectPoints);
        for (int j = 0; j < 4; j++)
        {
            line(result, rectPoints[j], rectPoints[(j + 1) % 4], Scalar(0, 255, 255), 2);
        }
        circle(result, rotatedRect.center, 5, Scalar(0, 0, 255), -1);
        string label = "L " + to_string(i + 1);
        putText(result, label, rotatedRect.center + Point2f(10, 0),
                FONT_HERSHEY_SIMPLEX, 0.5, Scalar(255, 255, 255), 1);
    }
    imshow("Bright (S>30)", mask); // 显示掩码
    return result;
}
int main()
{
    // 读取视频
    string videoPath = "/home/gerinskinut/Project/CExam/2024-autumn-selection/resource/test.avi";
    VideoCapture cap(videoPath);
    if (!cap.isOpened())
    {
        cout << "ERROR:Cannot open video file." << endl;
        return -1;
    }
    // 视频保存参数
    int frameWidth = static_cast<int>(cap.get(CAP_PROP_FRAME_WIDTH));
    int frameHeight = static_cast<int>(cap.get(CAP_PROP_FRAME_HEIGHT));
    double fps = cap.get(CAP_PROP_FPS);
    VideoWriter writer;
    string outputPath = "/home/gerinskinut/Project/CExam/2024-autumn-selection/PlanB/Output/marked_video.mp4";
    writer.open(outputPath, VideoWriter::fourcc('H', '2', '6', '4'), fps, Size(frameWidth, frameHeight));
    // 处理帧
    Mat frame, markedFrame;
    while (true)
    {
        cap >> frame;
        if (frame.empty())
        {
            break;
        }
        markedFrame = detectLightBars(frame);
        // 实时显示
        imshow("Light Bar Detection Result", markedFrame);
        if (writer.isOpened())
        {
            writer.write(markedFrame);
        }
        // 退出
        if (waitKey(1) == 'q')
        {
            break;
        }
    }
    // 释放
    cap.release();
    writer.release();
    destroyAllWindows();
    cout << "END." << outputPath << endl;
    return 0;
}